# auto_checklist.py
# usage: 比赛阶段正常不会使用这个程序，本程序用于通过samples文件自动标记checklist文件
# python ./src/tools/auto_checklist.py --checklist "./data/checklist.txt" --samples "./data/t_samples.round03.jsonl"

import argparse
import json
import sys
from typing import Dict, Tuple

def load_samples(samples_file: str) -> Dict[int, dict]:
    """加载 samples 文件，以 id 为键构建字典"""
    samples = {}
    try:
        with open(samples_file, 'r', encoding='utf-8') as f:
            for line_num, line in enumerate(f, start=1):
                line = line.strip()
                if not line:
                    continue
                try:
                    item = json.loads(line)
                    required_keys = {'id', 'question', 'sql', 'level'}
                    if not required_keys.issubset(item.keys()):
                        print(f"警告: 第 {line_num} 行在 --samples 中缺少必要字段: {item}")
                        continue
                    sample_id = int(item['id'])
                    samples[sample_id] = item
                except json.JSONDecodeError as e:
                    print(f"错误: 第 {line_num} 行 JSON 解析失败: {e}")
                    continue
    except FileNotFoundError:
        print(f"错误: 找不到文件 '{samples_file}'")
        sys.exit(1)
    except Exception as e:
        print(f"读取 samples 文件时发生错误: {e}")
        sys.exit(1)
    return samples

def main():
    parser = argparse.ArgumentParser(description="自动检查并更新 checklist 文件")
    parser.add_argument('--checklist', required=True, help='待处理的 checklist 文件路径 (JSONL)')
    parser.add_argument('--samples', required=True, help='样本数据文件路径 (JSONL)')
    args = parser.parse_args()

    # 加载 samples 数据
    samples = load_samples(args.samples)

    # 统计计数器
    count_question_mismatch = 0   # question 不一致
    count_sql_updated = 0         # sql 被更新
    count_matched = 0             # 完全匹配

    # 临时存储更新后的 checklist 内容
    updated_lines = []

    try:
        with open(args.checklist, 'r', encoding='utf-8') as f:
            lines = f.readlines()
    except FileNotFoundError:
        print(f"错误: 找不到文件 '{args.checklist}'")
        sys.exit(1)

    for line_num, line in enumerate(lines, start=1):
        line = line.strip()
        if not line:
            updated_lines.append(line + '\n')
            continue

        try:
            checklist_item = json.loads(line)
        except json.JSONDecodeError as e:
            print(f"第 {line_num} 行 JSON 解析失败: {e}")
            updated_lines.append(line + '\n')
            continue

        required_keys = {'id', 'question', 'sql', 'flag'}
        if not required_keys.issubset(checklist_item.keys()):
            print(f"第 {line_num} 行缺少必要字段，跳过: {checklist_item}")
            updated_lines.append(line + '\n')
            continue

        checklist_id = int(checklist_item['id'])

        # 查找 samples 中对应的记录
        if checklist_id not in samples:
            print(f"警告: id={checklist_id} 在 --samples 文件中未找到，跳过")
            updated_lines.append(line + '\n')
            continue

        sample_item = samples[checklist_id]

        # 比较 question
        if checklist_item['question'] != sample_item['question']:
            print(f"🚨 报警: id={checklist_id} 的 question 不一致!")
            print(f"   checklist: {checklist_item['question']}")
            print(f"   samples  : {sample_item['question']}")
            count_question_mismatch += 1
            updated_lines.append(line + '\n')  # 保持原样
            continue

        # question 一致，比较 sql
        if checklist_item['sql'] != sample_item['sql']:
            # 更新 sql 并设置 flag = -1
            checklist_item['sql'] = sample_item['sql']
            checklist_item['flag'] = -1
            count_sql_updated += 1
        else:
            # question 和 sql 都一致
            checklist_item['flag'] = 1
            count_matched += 1

        # 将更新后的行写入
        updated_lines.append(json.dumps(checklist_item, ensure_ascii=False) + '\n')

    # 写回更新后的 checklist 文件
    try:
        with open(args.checklist, 'w', encoding='utf-8') as f:
            f.writelines(updated_lines)
        print("\n✅ checklist 文件已更新。")
    except Exception as e:
        print(f"❌ 写入 checklist 文件失败: {e}")
        sys.exit(1)

    # 输出统计结果
    print("\n📊 处理结果统计:")
    print(f"  question 不一致（报警）: {count_question_mismatch}")
    print(f"  sql 已更新（flag=-1）  : {count_sql_updated}")
    print(f"  完全匹配（flag=1）     : {count_matched}")
    total_processed = count_question_mismatch + count_sql_updated + count_matched
    print(f"  总共处理条目: {total_processed}")

if __name__ == '__main__':
    main()